#include <jni.h>
#include <string.h>
#include <openssl/evp.h>
#include <openssl/ec.h>
#include <openssl/pem.h>
#include <openssl/err.h>

#include "sm2_crypto.h"

JNIEXPORT jint JNICALL Java_com_tianhe_common_tongsuo_SM2JNI_generateKeyPair(JNIEnv *env,
                                                   jobject obj,
                                                   jobjectArray keypair)
{
    EVP_PKEY *pstKeypair = NULL;
    BIO *pstPrivateBio = NULL;
    BIO *pstPublicBio = NULL;
    char *pcPemPrivateKey = NULL;
    char *pcPemPublicKey = NULL;
    long lPrivateKeyLen = 0;
    long lPublicKeyLen = 0;
    jstring jPrivateKey = NULL;
    jstring jPublicKey = NULL;

    if (sm2_keypair_generate(&pstKeypair) != 0) {
        fprintf(stderr, "Java_SM2JNI_generateKeyPair: sm2_keypair_generate failed.\n");
        return (-1);
    }

    pstPrivateBio = BIO_new(BIO_s_mem());
    pstPublicBio = BIO_new(BIO_s_mem());
    if (pstPrivateBio == NULL || pstPublicBio == NULL) {
        if (pstPrivateBio) {
            BIO_free(pstPrivateBio);
        }
        if (pstPublicBio) {
            BIO_free(pstPublicBio);
        }
        EVP_PKEY_free(pstKeypair);
        return (-1);
    }

    if (PEM_write_bio_PrivateKey(pstPrivateBio, pstKeypair, NULL, NULL, 0, NULL, NULL) == 0) {
        ERR_print_errors_fp(stderr);
        BIO_free(pstPrivateBio);
        BIO_free(pstPublicBio);
        EVP_PKEY_free(pstKeypair);
        return (-1);
    }

    if (PEM_write_bio_PUBKEY(pstPublicBio, pstKeypair) == 0) {
        ERR_print_errors_fp(stderr);
        BIO_free(pstPrivateBio);
        BIO_free(pstPublicBio);
        EVP_PKEY_free(pstKeypair);
        return (-1);
    }

    lPrivateKeyLen = BIO_get_mem_data(pstPrivateBio, &pcPemPrivateKey);
    lPublicKeyLen = BIO_get_mem_data(pstPublicBio, &pcPemPublicKey);
    if ((lPrivateKeyLen <= 0) || (lPublicKeyLen <= 0)){
        fprintf(stderr, "Java_SM2JNI_generateKeyPair: BIO_get_mem_data failed.\n");
        BIO_free(pstPrivateBio);
        BIO_free(pstPublicBio);
        EVP_PKEY_free(pstKeypair);
        return (-1);
    }

    jPrivateKey = (*env)->NewStringUTF(env, pcPemPrivateKey);
    jPublicKey = (*env)->NewStringUTF(env, pcPemPublicKey);
    if (jPrivateKey == NULL || jPublicKey == NULL) {
        fprintf(stderr, "Java_SM2JNI_generateKeyPair: create Java string failed.\n");
        BIO_free(pstPrivateBio);
        BIO_free(pstPublicBio);
        EVP_PKEY_free(pstKeypair);
        return (-1);
    }

    /* string array: [0] = private key, [1] = public key */
    (*env)->SetObjectArrayElement(env, keypair, 0, jPrivateKey);
    (*env)->SetObjectArrayElement(env, keypair, 1, jPublicKey);

    BIO_free(pstPrivateBio);
    BIO_free(pstPublicBio);
    EVP_PKEY_free(pstKeypair);

    return (0);
}

JNIEXPORT jint JNICALL Java_com_tianhe_common_tongsuo_SM2JNI_encrypt(JNIEnv *env,
                                           jobject obj,
                                           jstring pemPublicKey,
                                           jbyteArray id,
                                           jbyteArray plaintext,
                                           jbyteArray ciphertext)
{
    if (!pemPublicKey || !id || !plaintext || !ciphertext) {
        fprintf(stderr, "Java_SM2JNI_encrypt: invalid parameters.\n");
        return (-1);
    }

    const char *pcKey = (*env)->GetStringUTFChars(env, pemPublicKey, NULL);
    BIO *pstBio = BIO_new_mem_buf(pcKey, -1);
    if (pstBio == NULL) {
        fprintf(stderr, "Java_SM2JNI_encrypt: BIO_new_mem_buf failed.\n");
        return (-1);        
    }

    EVP_PKEY *pstPublicKey = PEM_read_bio_PUBKEY(pstBio, NULL, NULL, NULL);
    (*env)->ReleaseStringUTFChars(env, pemPublicKey, pcKey);
    BIO_free(pstBio);
    if (!pstPublicKey) {
        ERR_print_errors_fp(stderr);
        return -1;
    }

    jbyte *pucId = (*env)->GetByteArrayElements(env, id, NULL);
    jsize stIdLen = (*env)->GetArrayLength(env, id);
    jbyte *pucPlaintext = (*env)->GetByteArrayElements(env, plaintext, NULL);
    jsize stPlaintextLen = (*env)->GetArrayLength(env, plaintext);
    jbyte *pucCiphertext = (*env)->GetByteArrayElements(env, ciphertext, NULL);
    size_t stCiphertextLen = (*env)->GetArrayLength(env, ciphertext);
    jint iRet = -1;

    if (sm2_encrypt(pstPublicKey, (unsigned char *)pucId, stIdLen, (unsigned char *)pucPlaintext, stPlaintextLen, 
                    (unsigned char *)pucCiphertext, &stCiphertextLen) == -1) {
        fprintf(stderr, "Java_SM2JNI_encrypt: sm2_encrypt failed.\n");
    } else {
        iRet = (jint)stCiphertextLen;
    }

    (*env)->ReleaseByteArrayElements(env, id, pucId, JNI_ABORT);
    (*env)->ReleaseByteArrayElements(env, plaintext, pucPlaintext, JNI_ABORT);
    (*env)->ReleaseByteArrayElements(env, ciphertext, pucCiphertext, 0);
    EVP_PKEY_free(pstPublicKey);

    return (iRet);
}

JNIEXPORT jint JNICALL Java_com_tianhe_common_tongsuo_SM2JNI_decrypt(JNIEnv *env,
                                           jobject obj,
                                           jstring pemPrivateKey,
                                           jbyteArray id,
                                           jbyteArray ciphertext,
                                           jbyteArray plaintext) {
    if (!pemPrivateKey || !id || !ciphertext || !plaintext) {
        fprintf(stderr, "Java_SM2JNI_decrypt: invalid parameters.\n");
        return (-1);
    }

    const char *pcKey = (*env)->GetStringUTFChars(env, pemPrivateKey, NULL);
    BIO *pstBio = BIO_new_mem_buf(pcKey, -1);
    if (pstBio == NULL) {
        fprintf(stderr, "Java_SM2JNI_decrypt: BIO_new_mem_buf failed.\n");
        return (-1);        
    }

    EVP_PKEY *pstPrivateKey = PEM_read_bio_PrivateKey(pstBio, NULL, NULL, NULL);
    (*env)->ReleaseStringUTFChars(env, pemPrivateKey, pcKey);
    BIO_free(pstBio);

    if (!pstPrivateKey) {
        ERR_print_errors_fp(stderr);
        return -1;
    }

    jbyte *pucId = (*env)->GetByteArrayElements(env, id, NULL);
    jsize stIdLen = (*env)->GetArrayLength(env, id);
    jbyte *pucCiphertext = (*env)->GetByteArrayElements(env, ciphertext, NULL);
    jsize stCiphertextLen = (*env)->GetArrayLength(env, ciphertext);
    jbyte *pucPlaintext = (*env)->GetByteArrayElements(env, plaintext, NULL);
    size_t stPlaintextLen = (*env)->GetArrayLength(env, plaintext);
    jint iRet = -1;

    if (sm2_decrypt(pstPrivateKey, (unsigned char *)pucId, stIdLen, (unsigned char *)pucCiphertext, stCiphertextLen, 
                    (unsigned char *)pucPlaintext, &stPlaintextLen) == -1) {
        fprintf(stderr, "Java_SM2JNI_decrypt: sm2_decrypt failed.\n");
    } else {
        iRet = (jint)stPlaintextLen;
    }

    (*env)->ReleaseByteArrayElements(env, id, pucId, JNI_ABORT);
    (*env)->ReleaseByteArrayElements(env, ciphertext, pucCiphertext, JNI_ABORT);
    (*env)->ReleaseByteArrayElements(env, plaintext, pucPlaintext, 0);
    EVP_PKEY_free(pstPrivateKey);

    return (iRet);
}

JNIEXPORT jint JNICALL Java_com_tianhe_common_tongsuo_SM2JNI_sign(JNIEnv *env,
                                        jobject obj,
                                        jstring pemPrivateKey,
                                        jbyteArray id,
                                        jbyteArray msg,
                                        jbyteArray sign)
{
    if (!pemPrivateKey || !id || !msg || !sign) {
        fprintf(stderr, "Java_SM2JNI_sign: invalid parameters.\n");
        return (-1);
    }

    const char *pcKey = (*env)->GetStringUTFChars(env, pemPrivateKey, NULL);
    BIO *pstBio = BIO_new_mem_buf(pcKey, -1);
    if (pstBio == NULL) {
        fprintf(stderr, "Java_SM2JNI_sign: BIO_new_mem_buf failed.\n");
        return (-1);        
    }

    EVP_PKEY *pstPrivateKey = PEM_read_bio_PrivateKey(pstBio, NULL, NULL, NULL);
    (*env)->ReleaseStringUTFChars(env, pemPrivateKey, pcKey);
    BIO_free(pstBio);

    if (!pstPrivateKey) {
        ERR_print_errors_fp(stderr);
        return -1;
    }

    jbyte *pucId = (*env)->GetByteArrayElements(env, id, NULL);
    jsize stIdLen = (*env)->GetArrayLength(env, id);
    jbyte *pucMsg = (*env)->GetByteArrayElements(env, msg, NULL);
    jsize stMsgLen = (*env)->GetArrayLength(env, msg);
    jbyte *pucSign = (*env)->GetByteArrayElements(env, sign, NULL);
    size_t stSignLen = (*env)->GetArrayLength(env, sign);
    jint iRet = -1;

    if (sm2_sign(pstPrivateKey, (unsigned char *)pucId, stIdLen, (unsigned char *)pucMsg, stMsgLen,
                 (unsigned char *)pucSign, &stSignLen) == -1) {
        fprintf(stderr, "Java_SM2JNI_sign: sm2_sign failed.\n");
    } else {
        iRet = (jint)stSignLen;
    }

    (*env)->ReleaseByteArrayElements(env, id, pucId, JNI_ABORT);
    (*env)->ReleaseByteArrayElements(env, msg, pucMsg, JNI_ABORT);
    (*env)->ReleaseByteArrayElements(env, sign, pucSign, 0);
    EVP_PKEY_free(pstPrivateKey);

    return (iRet);
}

JNIEXPORT jint JNICALL Java_com_tianhe_common_tongsuo_SM2JNI_verify(JNIEnv *env,
                                          jobject obj,
                                          jstring pemPublicKey,
                                          jbyteArray id,
                                          jbyteArray msg,
                                          jbyteArray sign,
                                          jint signLen)
{
    if (!pemPublicKey || !id || !msg || !sign) {
        fprintf(stderr, "Java_SM2JNI_verify: invalid parameters.\n");
        return (-1);
    }

    const char *pcKey = (*env)->GetStringUTFChars(env, pemPublicKey, NULL);
    BIO *pstBio = BIO_new_mem_buf(pcKey, -1);
    if (pstBio == NULL) {
        fprintf(stderr, "Java_SM2JNI_verify: BIO_new_mem_buf failed.\n");
        return (-1);        
    }

    EVP_PKEY *pstPublicKey = PEM_read_bio_PUBKEY(pstBio, NULL, NULL, NULL);
    (*env)->ReleaseStringUTFChars(env, pemPublicKey, pcKey);
    BIO_free(pstBio);

    if (!pstPublicKey) {
        ERR_print_errors_fp(stderr);
        return -1;
    }

    jbyte *pucId = (*env)->GetByteArrayElements(env, id, NULL);
    jsize stIdLen = (*env)->GetArrayLength(env, id);
    jbyte *pucMsg = (*env)->GetByteArrayElements(env, msg, NULL);
    jsize stMsgLen = (*env)->GetArrayLength(env, msg);
    jbyte *pstSign = (*env)->GetByteArrayElements(env, sign, NULL);
    jint iRet = -1;

    iRet = sm2_verify(pstPublicKey, (unsigned char *)pucId, stIdLen, (unsigned char *)pucMsg, stMsgLen, (unsigned char *)pstSign,
                      signLen);
    if (iRet == -1) {
        fprintf(stderr, "Java_com_tianhe_common_tongsuo_SM2JNI_verify: sm2_verify failed.\n");
    }

    (*env)->ReleaseByteArrayElements(env, id, pucId, JNI_ABORT);
    (*env)->ReleaseByteArrayElements(env, msg, pucMsg, JNI_ABORT);
    (*env)->ReleaseByteArrayElements(env, sign, pstSign, JNI_ABORT);
    EVP_PKEY_free(pstPublicKey);

    return (iRet);
}